#!/usr/bin/env python3
# coding:utf-8
from argparse import ArgumentTypeError
from collections import defaultdict

import functools
import math
import threading
import csv
import time

import rclpy
from rclpy.node import Node
from rclpy.clock import Clock
from rclpy.clock import ClockType
from rclpy.qos import qos_profile_sensor_data
from ros2cli.node.direct import add_arguments as add_direct_node_arguments
from ros2cli.node.direct import DirectNode
from ros2topic.api import get_msg_class
from ros2topic.api import TopicNameCompleter
from ros2topic.verb import VerbExtension

DEFAULT_WINDOW_SIZE = 10000

# import debugpy
# debugpy.listen(6688)
# debugpy.wait_for_client()

Config = {
    "enable_record": True,
    "record_path": "/tmp/ros2_topic_hz",
    "csv_filed": [
        "time",
        "topic",
        "rate",
        "min",
        "max",
        "std_dev",
        "window",
    ],
}

sub_topics=[
    "/hz_string",
    "/hz_int",
    "/hz_int_array",
    "/hz_float",
    "/hz_float_array",
]


class ROSTopicHz(object):
    """ROSTopicHz receives messages for a topic and computes frequency."""

    def __init__(self, node, window_size, filter_expr=None, use_wtime=False):
        self.lock = threading.Lock()
        self.last_printed_tn = 0
        self.msg_t0 = -1
        self.msg_tn = 0
        self.times = []
        self._last_printed_tn = defaultdict(int)
        self._msg_t0 = defaultdict(lambda: -1)
        self._msg_tn = defaultdict(int)
        self._times = defaultdict(list)
        self.filter_expr = filter_expr
        self.use_wtime = use_wtime

        self.window_size = window_size

        # Clock that has support for ROS time.
        self._clock = node.get_clock()

    def get_last_printed_tn(self, topic=None):
        if topic is None:
            return self.last_printed_tn
        return self._last_printed_tn[topic]

    def set_last_printed_tn(self, value, topic=None):
        if topic is None:
            self.last_printed_tn = value
        self._last_printed_tn[topic] = value

    def get_msg_t0(self, topic=None):
        if topic is None:
            return self.msg_t0
        return self._msg_t0[topic]

    def set_msg_t0(self, value, topic=None):
        if topic is None:
            self.msg_t0 = value
        self._msg_t0[topic] = value

    def get_msg_tn(self, topic=None):
        if topic is None:
            return self.msg_tn
        return self._msg_tn[topic]

    def set_msg_tn(self, value, topic=None):
        if topic is None:
            self.msg_tn = value
        self._msg_tn[topic] = value

    def get_times(self, topic=None):
        if topic is None:
            return self.times
        return self._times[topic]

    def set_times(self, value, topic=None):
        if topic is None:
            self.times = value
        self._times[topic] = value

    def callback_hz(self, m, topic=None):
        """
        Calculate interval time.

        :param m: Message instance
        :param topic: Topic name
        """
        # ignore messages that don't match filter
        if self.filter_expr is not None and not self.filter_expr(m):
            return
        with self.lock:
            # Uses ROS time as the default time source and Walltime only if requested
            curr_rostime = self._clock.now() if not self.use_wtime else \
                Clock(clock_type=ClockType.SYSTEM_TIME).now()

            # time reset
            if curr_rostime.nanoseconds == 0:
                if len(self.get_times(topic=topic)) > 0:
                    print('time has reset, resetting counters')
                    self.set_times([], topic=topic)
                return

            curr = curr_rostime.nanoseconds
            msg_t0 = self.get_msg_t0(topic=topic)
            if msg_t0 < 0 or msg_t0 > curr:
                self.set_msg_t0(curr, topic=topic)
                self.set_msg_tn(curr, topic=topic)
                self.set_times([], topic=topic)
            else:
                self.get_times(topic=topic).append(curr - self.get_msg_tn(topic=topic))
                self.set_msg_tn(curr, topic=topic)

            if len(self.get_times(topic=topic)) > self.window_size:
                self.get_times(topic=topic).pop(0)

    def get_hz(self, topic=None):
        """
        Calculate the average publising rate.

        :param topic: topic name, ``list`` of ``str``
        :returns: tuple of stat results
            (rate, min_delta, max_delta, standard deviation, window number)
            None when waiting for the first message or there is no new one
        """
        if not self.get_times(topic=topic):
            return
        elif self.get_last_printed_tn(topic=topic) == 0:
            self.set_last_printed_tn(self.get_msg_tn(topic=topic), topic=topic)
            return
        elif self.get_msg_tn(topic=topic) < self.get_last_printed_tn(topic=topic) + 1e9:
            return
        with self.lock:
            # Get frequency every one minute
            times = self.get_times(topic=topic)
            n = len(times)
            mean = sum(times) / n
            rate = 1. / mean if mean > 0. else 0

            # std dev
            std_dev = math.sqrt(sum((x - mean)**2 for x in times) / n)

            # min and max
            max_delta = max(times)
            min_delta = min(times)

            self.set_last_printed_tn(self.get_msg_tn(topic=topic), topic=topic)

        return rate, min_delta, max_delta, std_dev, n

    def print_hz(self, topic=None):
        """Print the average publishing rate to screen."""
        ret = self.get_hz(topic)
        if ret is None:
            return
        rate, min_delta, max_delta, std_dev, window = ret
        print('%s, average rate: %.3f, min: %.3fs, max: %.3fs, std dev: %.5fs, window: %s'
              % (topic, rate * 1e9, min_delta * 1e-9, max_delta * 1e-9, std_dev * 1e-9, window))
        return

    def record_hz(self, topic, record_method):
        """Print the average publishing rate to screen."""
        ret = self.get_hz(topic)
        if ret is None:
            return
        rate, min_delta, max_delta, std_dev, window = ret
        cur_time = time.strftime("%Y%m%d_%H%M%S", time.localtime())
        rate = round(rate * 1e9,3)
        min_delta= round(min_delta * 1e-9,3)
        max_delta= round(max_delta * 1e-9,3)
        std_dev = round(std_dev * 1e-9,5)
        # data = f"{topic},{rate * 1e9:.2f},{min_delta * 1e-9:.3f},{max_delta * 1e-9:.3f},{std_dev * 1e-9:.3f},{window}"
        data = f"{topic},{rate},{min_delta},{max_delta},{std_dev},{window}"
        record_method(cur_time,topic,rate,min_delta, max_delta, std_dev, window)
        print('%s, average rate: %.3f, min: %.3fs, max: %.3fs, std dev: %.3fs, window: %s'
              % (topic, rate, min_delta, max_delta, std_dev, window))



class myTopic_Hz(Node):

    def __init__(self, sub_topics, config):
        super().__init__("my_topic_hz")
        self.config = config
        self.init_csv_record()
        self.rostopic_hz(sub_topics)

    def __exit__(self):
        if self.config["enable_record"] == True and self.csv_fp:
            self.csv_fp.close()

    def init_csv_record(self):
        if self.config["enable_record"] != True:
            self.csv_fp = None
            return
        time_stamp = time.strftime("%Y%m%d-%H%M%S", time.localtime())
        file = f"{self.config['record_path']}/hz_{time_stamp}.csv"
        fieldnames = self.config["csv_filed"]
        self.csv_fp = open(file, "w")
        self.csv_writer = csv.DictWriter(self.csv_fp, fieldnames=fieldnames)
        self.csv_writer.writeheader()

    def csv_record(self, cur_time, topic, rate, min_delta, max_delta, std_dev, window):
        if self.config["enable_record"] == False:
            return
        self.csv_writer.writerow(
            {
                "time": cur_time,
                "topic": topic,
                "rate": rate,
                "min": min_delta,
                "max": max_delta,
                "std_dev": std_dev,
                "window": window,
            }
        )

    def rostopic_hz(self, sub_topics, window_size=DEFAULT_WINDOW_SIZE, filter_expr=None, use_wtime=False):
        node = self
        sub_nodes = {}
        sub_msgs = []
        for topic in sub_topics:
            # pause hz until topic is published
            msg_class = get_msg_class(
                node, topic, blocking=True, include_hidden_topics=True)

            if msg_class is None:
                node.destroy_node()
                return

            rt = ROSTopicHz(node, window_size, filter_expr=filter_expr, use_wtime=use_wtime)
            node.create_subscription(
                msg_class,
                topic,
                functools.partial(rt.callback_hz, topic=topic),
                qos_profile_sensor_data)

            sub_msgs.append(msg_class)
            sub_nodes[topic] = rt

        enable_record = self.config["enable_record"]
        while rclpy.ok():
            rclpy.spin_once(node)
            for topic,rt in sub_nodes.items():
                if enable_record == True:
                    rt.record_hz(topic, self.csv_record)
                else:
                    rt.get_hz(topic)


        self.csv_fp.close()
        self.csv_fp=None
        node.destroy_node()
        rclpy.shutdown()



def main(args=None):
    rclpy.init(args=args)
    app = myTopic_Hz(sub_topics, Config)
    rclpy.spin(app)
    app.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()
